"""
annotation: todo
"""
import os
import shutil
import numpy as np
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset


def lerobot_builder():
    # Clean up any existing dataset in the output directory
    output_path = os.path.join(os.path.dirname(__file__), "../data/virtual_data")
    if os.path.exists(output_path):
        shutil.rmtree(output_path)

    # Create LeRobot dataset, define features to store
    # OpenPi assumes that proprio is stored in `state` and actions in `action`
    # LeRobot assumes that dtype of image data is `image`
    dataset = LeRobotDataset.create(
        repo_id='virtual_data',
        root=output_path,
        robot_type="panda",
        fps=10,
        features={
            "observation.image.image": {
                "dtype": 'video',
                "shape": (256, 256, 3),
                "names": ["height", "width", "channels"],
            },
            "observation.image.wrist_image": {
                "dtype": 'video',
                "shape": (256, 256, 3),
                "names": ["height", "width", "channels"],
            },
            "observation.state": {
                "dtype": "float32",
                "shape": (8,),
                "names": ["state"],
            },
            "action": {
                "dtype": "float32",
                "shape": (7,),
                "names": ["action"],
            },
        },
        image_writer_threads=10,
        image_writer_processes=5,
    )

    num_epoches = 6
    num_steps = 200
    for epoches in range(num_epoches):
        for step in range(num_steps + 10*epoches):
            task = f'task_{epoches}'
            if epoches == 0:
                task = f'task_{epoches}_step_{step}'
            dataset.add_frame(
                {
                    "observation.image.image": np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8),
                    "observation.image.wrist_image": np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8),
                    "observation.state": np.random.randn(8).astype(np.float32),
                    "actions": np.random.randn(7).astype(np.float32),
                    "task": task,
                }
            )
        dataset.save_episode()


if __name__ == "__main__":
    lerobot_builder()